import asyncio
import sys
import argparse
import signal
import re
from fastmcp import FastMCP
from fastmcp import Context
from fastmcp import settings
from mcp.server.lowlevel.server import LifespanResultT
from contextlib import (
    AbstractAsyncContextManager,
    AsyncExitStack,
    asynccontextmanager,
)
from collections.abc import AsyncIterator
from typing import Dict, Any
from client_manager import ClientManager
from connection_manager import ConnectionManager
from config import Config
from logger import Logger

config = Config()
logger = Logger("server")
client_manager = ClientManager(logger)

@asynccontextmanager
async def session_lifespan(app: FastMCP) -> AsyncIterator[Any]:
    """
    应用程序的生命周期管理上下文管理器。
    session创建时执行yield前的，session退出时执行yield后的。

    此上下文管理器在应用程序启动时初始化配置、日志记录、客户端管理器和支持提取标志。
    在应用程序关闭时，它会清理所有资源。

    Args:
        app: 正在运行的FastMCP应用实例

    Yields:
        一个空字典，作为应用程序的上下文对象
    """

#    logger.server_log("session连接：")
    try:
        yield {}
    finally:
        # 应用程序关闭时的清理操作
        # 停止扫描任务
        #logger.server_log("session退出")
        pass

# 创建新的MCP实例
mcp = FastMCP(name="Maipu AI Terminal", lifespan=session_lifespan)

@mcp.tool()
async def loggin(context: Context, user_id: str) -> Dict[str, Any]:
    """登录,不做认证,仅用于客户端管理
    参数:
    - user_id: 工号,仅限6位数字
    """
    if user_id is None or not re.match(r'^\d{6}$', user_id):
        return {
            'success': False,
            'message': f'{user_id} 格式错误, 必须是6位数字'
        }
    await client_manager.get_client_connection_manager_by_user_id(context.session_id, user_id)
    return {
        'success': True,
        'message': f'{user_id} 登录成功'
    }

@mcp.tool()
async def create_connection(context: Context, name: str, ip: str, port: int = 23) -> Dict[str, Any]:
    """连接设备
    参数:
    - name: 设备名称
    - ip: IP地址
    - port: 端口号，默认为23
    """
    try:
        retstr, client = await client_manager.get_conn_manager(context)
        if not client:
            return {
                'success': False,
                'message': retstr
            }
        connection = await client.create_connection(name, ip, port)
        
        # 检查connection_manager返回的是否是错误信息
        if isinstance(connection, dict) and 'error' in connection:
            client.logger.server_log(f"创建连接 {name} 失败: {connection['error']}")
            return {
                'success': False,
                'message': connection['error'],
                'error_code': 'CONNECTION_CREATE_FAILED'
            }
        
        return {
            'success': True,
            'message': f'连接 {name} 已创建',
            'connection': {
                'name': name,
                'ip': ip,
                'port': port
            }
        }
    except Exception as e:
        logger.server_log(f'创建连接 {name} 失败: {str(e)}')
        return {
            'success': False,
            'message': f'创建连接失败: {str(e)}'
        }

@mcp.tool()
async def batch_create_connections(context: Context, connections: list) -> Dict[str, Any]:
    """批量连接设备
    参数格式:
    [
        {'name': '设备名1', 'ip': 'IP地址1', 'port': 端口号1},
        {'name': '设备名2', 'ip': 'IP地址2', 'port': 端口号2},
        ...
    ]
    """
    try:
        retstr, client = await client_manager.get_conn_manager(context)
        if not client:
            return {
                'success': False,
                'message': retstr
            }
        results = []
        success_count = 0
        
        for conn_info in connections:
            name = conn_info.get('name')
            ip = conn_info.get('ip')
            port = conn_info.get('port', 23)  # 默认端口23
            
            if not name or not ip:
                result = {
                    'success': False,
                    'message': '连接名称和IP地址是必需的',
                    'connection': conn_info
                }
                results.append(result)
                continue
            
            try:
                connection = await client.create_connection(name, ip, port)
                        
                # 检查connection_manager返回的是否是错误信息
                if isinstance(connection, dict) and 'error' in connection:
                    client.logger.server_log(f"创建连接 {name} 失败: {connection['error']}")
                    result = {
                        'success': False,
                        'message': connection['error']
                    }
                else:
                    result = {
                        'success': True,
                        'message': f'连接 {name} 已创建',
                        'connection': {
                            'name': name,
                            'ip': ip,
                            'port': port
                        }
                    }
                    success_count += 1
            except Exception as e:
                result = {
                    'success': False,
                    'message': f'创建连接 {name} 失败: {str(e)}',
                    'connection': conn_info
                }
            
            results.append(result)
        
        return {
            'success': True,
            'message': f'批量创建完成，成功 {success_count}/{len(connections)} 个连接',
            'results': results
        }
    except Exception as e:
        client.logger.server_log(f'批量创建连接失败: {str(e)}')
        return {
            'success': False,
            'message': f'批量创建连接失败: {str(e)}'
        }

@mcp.tool()
async def close_connection(context: Context, name: str) -> Dict[str, Any]:
    """取消连接设备
    参数:
    - name: 要取消连接的设备名称
    """

    try:
        retstr, client = await client_manager.get_conn_manager(context)
        if not client:
            return {
                'success': False,
                'message': retstr
            }
        ret = await client.close_connection(name)
        # 检查connection_manager返回的是否是错误信息
        if isinstance(ret, dict) and 'error' in ret:
            client.logger.server_log(f"关闭连接 {name} 失败: {ret['error']}")
            return {
                'success': False,
                'message': ret['error']
            }
        return {
            'success': True,
            'message': f'连接 {name} 已关闭'
        }
    except Exception as e:
        client.logger.server_log(f'关闭连接 {name} 失败: {str(e)}')
        return {
            'success': False,
            'message': f'关闭连接失败: {str(e)}'
        }

@mcp.tool()
async def list_connections(context: Context) -> Dict[str, Any]:
    """查看所有已连接的设备
    """

    try:
        retstr, client = await client_manager.get_conn_manager(context)
        if not client:
            return {
                'success': False,
                'message': retstr
            }
        connections = await client.list_connections()
        return {
            'success': True,
            'connections': connections
        }
    except Exception as e:
        client.logger.server_log(f'列出连接失败: {str(e)}')
        return {
            'success': False,
            'message': f'列出连接失败: {str(e)}'
        }

@mcp.tool()
async def send_command(context: Context, name: str, command: str) -> Dict[str, Any]:
    """向指定设备发送命令
    参数:
    - name: 设备名称
    - command: 要发送的命令
    """
    try:
        retstr, client = await client_manager.get_conn_manager(context)
        if not client:
            return {
                'success': False,
                'message': retstr
            }
        result = await client.send_command(name, command)
        return {
            'success': True,
            'result': result
        }
    except Exception as e:
        client.logger.server_log(f'发送命令到连接 {name} 失败: {str(e)}')
        return {
            'success': False,
            'message': f'发送命令失败: {str(e)}'
        }


async def set_log_dir(context: Context, log_dir: str) -> Dict[str, Any]:
    """设置日志目录
    参数:
    - log_dir: 日志文件夹路径
    """
    try:
        config.set_log_dir(log_dir)
        retstr, client = await client_manager.get_conn_manager(context)
        if not client:
            return {
                'success': False,
                'message': retstr
            }
        client.logger.set_log_dir(log_dir)
        client.logger.server_log(f'日志目录已设置为: {log_dir}')
        return {
            'success': True,
            'message': f'日志目录已设置为: {log_dir}'
        }
    except Exception as e:
        client.logger.server_log(f'设置日志目录失败: {str(e)}')
        return {
            'success': False,
            'message': f'设置日志目录失败: {str(e)}'
        }

@mcp.tool()
async def close_all_connections(context: Context) -> Dict[str, Any]:
    """取消连接所有设备
    """
    try:
        retstr, client = await client_manager.get_conn_manager(context)
        if not client:
            return {
                'success': False,
                'message': retstr
            }
        count = await client.close_all_connections()
        client.logger.server_log(f'已成功关闭会话 {context.session_id} 的 {count} 个连接')
        return {
            'success': True,
            'message': f'已成功关闭 {count} 个连接',
            'closed_count': count
        }
    except Exception as e:
        client.logger.server_log(f'关闭所有连接失败: {str(e)}')
        return {
            'success': False,
            'message': f'关闭所有连接失败: {str(e)}'
        }

@mcp.tool()
async def broadcast_command(context: Context, command: str) -> Dict[str, Any]:
    """向所有设备广播发送命令
    参数:
    - command: 要发送的命令
    返回每个设备的命令执行结果
    """
    try:
        retstr, client = await client_manager.get_conn_manager(context)
        if not client:
            return {
                'success': False,
                'message': retstr
            }
        results = await client.broadcast_command(command)
        return {
            'success': True,
            'results': results
        }
    except Exception as e:
        client.logger.server_log(f'广播命令失败: {str(e)}')
        return {
            'success': False,
            'message': f'广播命令失败: {str(e)}'
        }

def setup_argparse():
    parser = argparse.ArgumentParser(description='Maipu AI Terminal - MCP服务器')

    parser.add_argument('--log_dir', default='./logs', help='日志文件夹路径')
    parser.add_argument('--transport', default='sse', help='传输方式，"stdio", "sse", or "streamable-http"，默认sse')
    parser.add_argument('--host', default='0.0.0.0', help='服务器主机地址，默认0.0.0.0')
    parser.add_argument('--port', type=int, default=18888, help='服务器端口，默认18888')

    return parser

def init(args):
    # 初始化主线程配置和日志
    config.set_log_dir(args.log_dir)
    logger.set_log_dir(args.log_dir)
    
    # 注册连接并打印结果
    result = logger.register_connection('server', 'localhost', 0)
    
    # 记录服务器日志
    logger.server_log('MCP服务器线程已启动')
    logger.server_log(f'日志目录: {config.log_dir}')
    logger.server_log(f'传输方式: {args.transport}')

def main():
    parser = setup_argparse()
    args = parser.parse_args()

    init(args)
    
    # 运行服务器
    try:
        if args.transport=="stdio":
            mcp.run(transport=args.transport)
        else:
            mcp.run(transport=args.transport, host=args.host, port=args.port)
    finally:
        client_manager.stop_scanner()
        client_manager.close_all_clients()

if __name__ == '__main__':
    main()
